/*
 * Copyright The OpenTelemetry Authors
 * SPDX-License-Identifier: Apache-2.0
 */

package com.xiashitech.agent.instrumentation.httpurlconnection.inst;

import com.xiashitech.agent.cache.status.CacheStatus;
import com.xiashitech.agent.instrumentation.base.propagator.W3CReentryBaggagePropagator;
import com.xiashitech.agent.instrumentation.bridge.Java8Bridge;
import com.xiashitech.agent.instrumentation.helper.calldepth.XSHttpURLConnectionCallDepth;
import com.xiashitech.agent.instrumentation.httpurlconnection.common.ConnectionHeaderGetter;
import com.xiashitech.agent.utils.DurationPrinter;
import io.opentelemetry.javaagent.bootstrap.CallDepth;
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;
import net.bytebuddy.matcher.ElementMatchers;

import java.net.HttpURLConnection;

import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.extendsClass;
import static net.bytebuddy.matcher.ElementMatchers.*;

public class XSHttpUrlConnectionInstrumentation implements TypeInstrumentation {
  @Override
  public ElementMatcher<TypeDescription> typeMatcher() {
    return nameStartsWith("java.net.")
        .or(ElementMatchers.<TypeDescription>nameStartsWith("sun.net"))
        .or(named("weblogic.net.http.HttpURLConnection"))
        .and(not(named("sun.net.www.protocol.https.HttpsURLConnectionImpl")))
        .and(extendsClass(named("java.net.HttpURLConnection")));
  }

  @Override
  public void transform(TypeTransformer transformer) {
    transformer.applyAdviceToMethod(
        isMethod().and(isPublic()).and(named("getInputStream")),
        this.getClass().getName() + "$HttpUrlConnectionAdvice");
  }

  @SuppressWarnings("unused")
  public static class HttpUrlConnectionAdvice {

    @Advice.OnMethodEnter(suppress = Throwable.class)
    public static void methodEnter(
            @Advice.Local("otelCallDepth") CallDepth callDepth) {
      if(CacheStatus.cannotRun()) {
        return;
      }
      callDepth = CallDepth.forClass(XSHttpURLConnectionCallDepth.class);
      if (callDepth.getAndIncrement() > 0) {
        return;
      }
    }

    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
    public static void methodExit(
            @Advice.This HttpURLConnection connection,
            @Advice.Local("otelCallDepth") CallDepth callDepth) {
      DurationPrinter durationPrinter = new DurationPrinter("XSHttpUrlConnectionInstrumentation");
      if(CacheStatus.cannotRun()) {
        return;
      }
      if (callDepth.decrementAndGet() > 0) {
        return;
      }
      // prevent infinite recursion in case end() captures response headers due to
      // HttpUrlConnection.getHeaderField() calling HttpUrlConnection.getInputStream() which then
      // enters this advice again
      callDepth.getAndIncrement();
      try {
        W3CReentryBaggagePropagator.getInstance().extract(Java8Bridge.currentForXSContext(), connection, ConnectionHeaderGetter.INSTANCE).makeCurrent();
      } finally {
        callDepth.decrementAndGet();
      }
      durationPrinter.printElapseDuration("HttpUrlConnectionAdvice e");
    }
  }
}
